package com.wisecoach.gatewayplus.transaction.lockedtx;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ValueOperations;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * {@code @author:} wisecoach
 * {@code @date:} 2023/9/30 下午7:21
 * {@code @version:} 1.0.0
 */


public class RedisLockSource implements LockSource {

    private final RedisTemplate<String, Map<String, RWLock>> redisTemplate;

    public RedisLockSource(RedisTemplate<String, Map<String, RWLock>> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public boolean lock(TxRWSet rwSet) {
        redisTemplate.multi();
        ValueOperations<String, Map<String, RWLock>> operations = redisTemplate.opsForValue();
        AtomicBoolean allLocked = new AtomicBoolean();
        rwSet.getWriteSet().forEach((ns, keys) -> {
            operations.setIfAbsent(ns, new HashMap<>());
            Map<String, RWLock> nsMap = operations.get(ns);
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                if (!rwLock.canLock()) {
                    allLocked.compareAndSet(true, false);
                }
            });
        });
        rwSet.getReadSet().forEach((ns, keys) -> {
            operations.setIfAbsent(ns, new HashMap<>());
            Map<String, RWLock> nsMap = operations.get(ns);
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                if (!rwLock.canRLock()) {
                    allLocked.compareAndSet(true, false);
                }
            });
        });
        if (!allLocked.get()) {
            redisTemplate.discard();
            return allLocked.get();
        }
        rwSet.getWriteSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = operations.get(ns);
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                boolean tryLock = rwLock.tryLock();
                if (!tryLock) {
                    throw new RuntimeException("出了些意外");
                }
            });
            operations.set(ns, nsMap);
        });
        rwSet.getReadSet().forEach((ns, keys) -> {
            Map<String, RWLock> nsMap = operations.get(ns);
            keys.forEach(key -> {
                RWLock rwLock = nsMap.computeIfAbsent(key, k -> RWLock.newRWLock());
                boolean tryLock = rwLock.tryRLock();
                if (!tryLock) {
                    throw new RuntimeException("出了些意外");
                }
            });
            operations.set(ns, nsMap);
        });
        redisTemplate.exec();
        return allLocked.get();
    }

    @Override
    public void unLock(TxRWSet rwSet) {
        redisTemplate.multi();
        ValueOperations<String, Map<String, RWLock>> operations = redisTemplate.opsForValue();
        try {
            rwSet.getWriteSet().forEach((ns, keys) -> {
                Map<String, RWLock> nsMap = operations.get(ns);
                keys.forEach(key -> nsMap.get(key).unlock());
                operations.set(ns, nsMap);
            });
            rwSet.getReadSet().forEach((ns, keys) -> {
                Map<String, RWLock> nsMap = operations.get(ns);
                keys.forEach(key -> nsMap.get(key).unRLock());
                operations.set(ns, nsMap);
            });
            redisTemplate.exec();
        } catch (NullPointerException e) {
            redisTemplate.discard();
            throw new RuntimeException(e);
        }
    }

    private Map<String, Set<String>> getReadSetToLock(TxRWSet rwSet) {
        HashMap<String, Set<String>> ret = new HashMap<>(rwSet.getReadSet());
        rwSet.getWriteSet().keySet().forEach(ret::remove);
        return ret;
    }

}
